/* Based on Density1D.c
Estimates density of individuals, mean trait values, and trait variances at each individual's location
Assumes a Gaussian kernel as the neighbourhood size.
Requires an array of 1D locations (X[i]), trait values (H[i], D[i])
bandwidth currently fixed at bw=1 (easy enough to make variable)
*/

#include <R.h>
#include <Rinternals.h>
//#include <Rmath.h>
//#include <Rdefines.h>
#include <math.h>
#include <omp.h> // parallels library
#include "Rinterface.h"

#define CHUNKSIZE 1 /*defines the chunk size as 1 contiguous iteration*/
#define Pi 3.141593
#define CSTACK_DEFNS 7


//--------------------------------------//
// Function declarations //

double norm (double x, double bw);
SEXP metrics (SEXP R_X, SEXP R_H, SEXP D_H, SEXP R_n, SEXP R_ev);
SEXP pmetrics (SEXP R_X, SEXP R_H, SEXP R_D, SEXP R_n, SEXP R_ev, SEXP R_ncores);
SEXP sum_metrics (SEXP R_X, SEXP R_H, SEXP R_D, SEXP R_n, SEXP R_bins, SEXP R_nbins, SEXP R_ev, SEXP R_bw);


//--------------------------------------//
//  Function definitions  //

double norm (double x, double bw){
	return exp(-pow(x,2)/(2*pow(bw,2)))/(bw*sqrt(2*Pi));
}

SEXP metrics (SEXP R_X, SEXP R_H, SEXP R_D, SEXP R_n, SEXP R_ev){
	int n = INTEGER(coerceVector(R_n, INTSXP))[0]; //grab vector length
	int ev = INTEGER(coerceVector(R_ev, INTSXP))[0];
	R_X=coerceVector(R_X, REALSXP); //digest R_X
	R_H=coerceVector(R_H, REALSXP); //digest R_H
	R_D=coerceVector(R_D, REALSXP); //digest R_D
	SEXP outmat; PROTECT(outmat=allocMatrix(REALSXP, n, 5)); // a place to put all the output
	/*
	SEXP dens; PROTECT(dens = allocVector(REALSXP,n)); //a place to put the density array
	SEXP hmean; PROTECT(hmean = allocVector(REALSXP, n)); // and the hmean array
	SEXP dmean; PROTECT(dmean = allocVector(REALSXP, n)); // and the dmean array
	SEXP hvar; PROTECT(hvar = allocVector(REALSXP, n)); // and the hvar array
	SEXP dvar; PROTECT(dvar = allocVector(REALSXP, n)); // and the dvar array
	*/
	double *X, *RH, *RD, *out;  //pointers variables
	double w; //to take weights
	
	X = REAL(R_X); //pointers to real parts of R vectors
	RH = REAL(R_H);
	RD = REAL(R_D);
	out = REAL(outmat);
	//N = REAL(dens); //pointer to output

	//D = REAL(dmean);
	//H = REAL(hmean);
	//Dv = REAL(dvar);
	//Hv = REAL(hvar);
	
	//calculate density and mean trait values
	for (int ii=0; ii<n; ii++){
		out[ii] = 0; //column 1 for density
		out[ii+n] = 0; //column 2 for meanD
		out[ii+2*n] = 0; //column 3 for meanH
		for (int jj=0; jj<n; jj++){
			w = norm(X[ii]-X[jj], 1);
			out[ii] += w;
			out[ii+n] += w*(RD[jj]);
			out[ii+2*n] += w*(RH[jj]);
		}
		out[ii+n] = out[ii+n]/out[ii];
		out[ii+2*n] = out[ii+2*n]/out[ii];
	}
	
	//calculate trait variances
	if (ev==1){
		for (int ii=0; ii<n; ii++){
			out[ii+3*n] = 0;
			out[ii+4*n] = 0;
			for (int jj=0; jj<n; jj++){
				w = norm(X[ii]-X[jj], 1);
				out[ii+3*n] += w*pow(out[ii+n]-RD[jj], 2);
				out[ii+4*n] += w*pow(out[ii+2*n]-RH[jj], 2);
			}
			out[ii+3*n] = sqrt(out[ii+3*n]/out[ii]);
			out[ii+4*n] = sqrt(out[ii+4*n]/out[ii]);
		}
	}
	
	UNPROTECT(1);
	return(outmat);
}


// A parallel version of the metrics function
SEXP pmetrics (SEXP R_X, SEXP R_H, SEXP R_D, SEXP R_n, SEXP R_ev, SEXP R_ncores){
	//R_CStackLimit=(uintptr_t)-1; // trick R into coping with threads
	int n = INTEGER(coerceVector(R_n, INTSXP))[0]; //grab vector length
	int ev = INTEGER(coerceVector(R_ev, INTSXP))[0];
    int ncores = INTEGER(coerceVector(R_ncores, INTSXP))[0];
	R_X=coerceVector(R_X, REALSXP); //digest R_X
	R_H=coerceVector(R_H, REALSXP); //digest R_H
	R_D=coerceVector(R_D, REALSXP); //digest R_D
	SEXP outmat; PROTECT(outmat=allocMatrix(REALSXP, n, 5)); // a place to put all the output
	
	double *X, *RH, *RD, *out;  //pointers variables
	double w; //to take weights
    int ii, jj; // iterations variables
	
	X = REAL(R_X); //pointers to real parts of R vectors
	RH = REAL(R_H);
	RD = REAL(R_D);
	out = REAL(outmat);
		


    omp_set_num_threads(ncores);   /* requests nthread threads */
    printf("nprocs = %i -- nthreads = %i\n", omp_get_num_procs(), omp_get_max_threads());
    #pragma omp parallel private(ii, jj, w)
    {
        #pragma omp for schedule(dynamic, CHUNKSIZE)  //only make outer loop parallel
        //calculate density and mean trait values
        for (ii=0; ii<n; ii++){
            out[ii] = 0; //column 1 for density
            out[ii+n] = 0; //column 2 for meanD
            out[ii+2*n] = 0; //column 3 for meanH
            for (jj=0; jj<n; jj++){
                w = norm(X[ii]-X[jj], 1);
                out[ii] += w;
                out[ii+n] += w*(RD[jj]);
                out[ii+2*n] += w*(RH[jj]);
            }
            out[ii+n] = out[ii+n]/out[ii];
            out[ii+2*n] = out[ii+2*n]/out[ii];
        }
    
     }
	//calculate trait variances
	if (ev==1){
		#pragma omp parallel private(ii, jj, w)
		{	
			#pragma omp for schedule(dynamic, CHUNKSIZE)
			for (ii=0; ii<n; ii++){
				out[ii+3*n] = 0;
				out[ii+4*n] = 0;
				for (jj=0; jj<n; jj++){
					w = norm(X[ii]-X[jj], 1);
					out[ii+3*n] += w*pow(out[ii+n]-RD[jj], 2);
					out[ii+4*n] += w*pow(out[ii+2*n]-RH[jj], 2);
				}
				out[ii+3*n] = sqrt(out[ii+3*n]/out[ii]);
				out[ii+4*n] = sqrt(out[ii+4*n]/out[ii]);
			}
		}
	}
	
	UNPROTECT(1);
	return(outmat);
}


// A serial version of the metrics function that calculates metrics back to fixed points rather than individuals
SEXP sum_metrics (SEXP R_X, SEXP R_H, SEXP R_D, SEXP R_n, SEXP R_bins, SEXP R_nbins, SEXP R_ev, SEXP R_bw){
	int n = INTEGER(coerceVector(R_n, INTSXP))[0]; //grab vector length
	int nbins = INTEGER(coerceVector(R_nbins, INTSXP))[0]; //grab bin length
	int ev = INTEGER(coerceVector(R_ev, INTSXP))[0];
	int bw = INTEGER(coerceVector(R_bw, INTSXP))[0];
	
	R_X=coerceVector(R_X, REALSXP); //digest R_X
	R_H=coerceVector(R_H, REALSXP); //digest R_H
	R_D=coerceVector(R_D, REALSXP); //digest R_D
	R_bins=coerceVector(R_bins, INTSXP);
	SEXP outmat; PROTECT(outmat=allocMatrix(REALSXP, nbins, 5)); // a place to put all the output
	
	double *X, *RH, *RD, *out;  //pointers variables
	double w; //to take weights
    int *Rb, ii, jj; // iterations variables
	
	X = REAL(R_X); //pointers to real parts of R vectors
	RH = REAL(R_H);
	RD = REAL(R_D);
	out = REAL(outmat);
	Rb = INTEGER(R_bins);
		
	//calculate density and mean trait values
	for (ii=0; ii<nbins; ii++){
		out[ii] = 0; //column 1 for density
		out[ii+nbins] = 0; //column 2 for meanD
		out[ii+2*nbins] = 0; //column 3 for meanH
		for (jj=0; jj<n; jj++){
			w = norm(Rb[ii]-X[jj], bw);
			out[ii] += w;
			out[ii+nbins] += w*(RD[jj]);
			out[ii+2*nbins] += w*(RH[jj]);
		}
		out[ii+nbins] = out[ii+nbins]/out[ii];
		out[ii+2*nbins] = out[ii+2*nbins]/out[ii];
	}
	
	//calculate trait variances
	if (ev==1){
			for (ii=0; ii<nbins; ii++){
				out[ii+3*nbins] = 0;
				out[ii+4*nbins] = 0;
				for (jj=0; jj<n; jj++){
					w = norm(Rb[ii]-X[jj], bw);
					out[ii+3*nbins] += w*pow(out[ii+nbins]-RD[jj], 2);
					out[ii+4*nbins] += w*pow(out[ii+2*nbins]-RH[jj], 2);
				}
				out[ii+3*nbins] = sqrt(out[ii+3*nbins]/out[ii]);
				out[ii+4*nbins] = sqrt(out[ii+4*nbins]/out[ii]);
			}
	}
	
	UNPROTECT(1);
	return(outmat);
}


